from enthought.traits.api import HasTraits,Int,Float,Str,Property,Range,Array, Instance
from enthought.traits.ui.api import View,Item,Label
from enthought.chaco.chaco_plot_editor import ChacoPlotItem
from chaco.api import Plot, ArrayPlotData
from numpy import array,poly1d,arange,linspace
from enable.component_editor import ComponentEditor
from chaco import default_colors as dc
import numpy as np
import filereader as gds

class WaterfallPlotter(HasTraits):
    '''Creates a waterfall plot from 2D datasets.'''

    # Creates some sample data for testing.
    xslider = Range(0, 20)
    yslider = Range(25000, 75000)

    xoffset = Property(Float, depends_on=["xslider"])
    yoffset = Property(Float, depends_on=["yslider"])
    
    plot = Instance(Plot)

    # Sets default values.
    def _intenseslider_default(self): return 50000
    def _thetaslider_default(self): return 50000

    def _xdata_default(self): return linspace(0, 10, 10)

    def _get_ydata(self): return self.slider * self.xdata 
    traits_view = View(
             Item('plot',editor=ComponentEditor(), show_label=False),
             width=900, height=600, resizable=True, title="Waterfall Plotter")

    def __init__(self, data):
        super(WaterfallPlotter, self).__init__()
        plotdata = ArrayPlotData()
        plot = Plot(plotdata)

        extraX = 0
        extraY = 0
        cr = 0
        cp = 0
        
        for key in data: 
            xlabel = "x%d" % key
            ylabel = "y%d" % key
            extraX += 0.3
            extraY += 3000
            cr += 1
            if cr == 10:
                cr = 0
                cp += 1
            if cp == 3:
                cp = 0
            
            data[key][0] += extraX
            data[key][1] += extraY

            plotdata.set_data(xlabel, data[key][0])
            plotdata.set_data(ylabel, data[key][1])
            plot.plot((xlabel, ylabel), type="line", color=tuple(dc.PALETTES[cp][cr]))
        
        plot.title = "Title Here"
        self.plot = plot

        '''    traits_view = View(
                ChacoPlotItem("xdata","ydata", y_auto=True,resizable=True,
                              show_label=False, x_label="2Theta",
                              y_label="Intensity", title="Title Here"),
                Item('slider'),
                resizable = True, 
                width=900, height=600,               
                title="Waterfall Plotter"               
                      )
        '''

if __name__ == "__main__":
    '''   fdata = np.loadtxt("LiRh2O4_c2k_083K-00312.chi", skiprows = 4)
    fdata = np.transpose(fdata)
    data = np.array([fdata])
    
    fdata = np.loadtxt("LiRh2O4_c2k_087K-00310.chi", skiprows = 4)
    extraX = 50
    extraY = 50000
    fdata[:,0] += extraX
    fdata[:,1] += extraY
    fdata = np.transpose(fdata)

    data = np.append(data, [fdata], axis=0)
    '''
    data = gds.getDataSets('.chi')
    p = WaterfallPlotter(data)
    p.configure_traits()

